"""
File: big_csv
Author: Liangxu.CHEN
Data: 2024/4/29
"""
import os
import pandas as pd
# 当前文件夹路径
current_directory = os.path.dirname(os.path.abspath(__file__))


def get_tasks(search_key):
    # 创建一个新的Dataframe数据
    res_df = pd.DataFrame()
    print("正在搜索文件夹:", current_directory)
    for file in os.listdir(current_directory):
        if file.split(".")[-1] == "csv":
            data_file_path = os.path.join(current_directory, file)
            step = 5e4  # 科学记数法 50000
            max_row = 1e6  # 一百万
            row_number = 1  # 本文档第一行没有用途，直接跳过
            try:
                while row_number < max_row:
                    # 读取文件，由于文件较大，skiprows代表忽略行数，nrows可以指定获取数据行数
                    df = pd.read_csv(data_file_path, encoding="gbk", skiprows=row_number, nrows=row_number + step)
                    df = df[df['搜索词'].str.contains(search_key)]
                    # 合并数据
                    res_df = pd.concat([res_df, df])
                    row_number += step
            except Exception as e:
                # 有错误，跳出循环，但不要影响后续的操作
                print("运行错误：", e)
                break
    return res_df


if __name__ == '__main__':
    # 创建结果文件夹，保存所有提取的数据电子表格
    path = os.path.join(current_directory, "结果")
    # 若没有这个文件夹则创建新的
    if not os.path.exists(path):
        os.mkdir(path)
    key_words = input("输入你的关键词(多个关键词可用‘,’分割):")
    key_words = key_words.split(",")
    for index, key in enumerate(key_words):
        print("正在搜索关键词：", key)
        result_df = get_tasks(key)
        # 每个关键词保存一个新的电子表格
        file_path = os.path.join(path, f"{key}.csv")
        result_df.to_csv(file_path)